import math
import numpy
import pylab
import re
import copy

import operator


#colormap..?
#redMap = pylab.matplotlib.colors.LinearSegmentedColormap('redMap',
#							 {'red':   [(0.0,0.0,0.0),
#								    (1.0,1.0,1.0)],
#							  'green': [(0.0,0.0,0.0),
#								    (1.0,0.0,0.0)],
#							  'blue':  [(0.0,0.0,0.0),
#								    (1.0,1.0,1.0)]}
#							 )



class Movie:

    ''' represents a movie.
    '''

    def __init__(self, activityMatrix, slicePtr, movieID, shortID=None):
	'''
	@param matrix : list of lists where each sublist a neuron's binary spike activity per time bin.
	@param slicePtr : this is important because two movies will only operate on eachother if they're from the same slice.
	'''
	#make a true 2-dimensional matrix out of it:
	self.data = numpy.array(activityMatrix)

	self.ID = movieID
	if(shortID==None):
	    self.shortID = self.ID
	else:
	    self.shortID=shortID
	##eventually i want to give movies a shortID, like just the number,
	##to make it easier to name movies that are formed from operations on
        ##several movies.
		
	#the way i get at this via self.data is a little roundabout, 
	#but if we only ever reference self.data, 
	#then the valid types for activityMatrix are anything the numpy.array constructor
	#accepts -- type flexibility for free!

	#how many neurons
	self.numNeurons = len(self.data) #this returns the number of rows
	#how many frames:
	#data.size is the total number of entries in the matrix, so dividing by numRows gives num Columns
	self.numBins = self.data.size / self.numNeurons 

	self.parentSlice = slicePtr

	#make a 1xnumNeurons array that's 1 for any neuron that was ever on.
	#we have a special step for ComboMovies
	print(repr(self))
	if(repr(self).find('Combo') >= 0):
	    tdata = map(lambda y : numpy.int32(numpy.ceil(y)), self.data)
	    self.timeless = [reduce(operator.or_, neuron) for neuron in tdata]
	else:
	    self.timeless = [reduce(operator.or_, neuron) for neuron in activityMatrix]
	      

	#you can enter a protocol for each movie after you read them in. for now.
	self.protocol = ''


#-----------------------------------------------------------------------------------------------naiveAnd
    def naiveAnd(self, other):
	'''
	performs a frame-by-frame AND of the two movies until one or the other is out of frames 
	(no alignment - just ANDs the first frames, then the second ...)
	returns a new Movie instance.
	'''
	self._checkCompat(other)
	#determine how many frames the new movie will have:
	numFrames = min(self.numBins, other.numBins)
	#the .data field of each instance is neuronsXtimesteps, 
	#but because python doesn't provide the illusion of true 2D arrays, (it's represented as an array of rows)
	#this operation is more direct if we do it on the transpose , and then transpose the result.
	temp = numpy.array([map(operator.__and__, self.data.T[i], other.data.T[i]) for i in range(0,numFrames)])
	return Movie(temp.T, self.parentSlice, ''.join(['(', self.ID, ' naiveAnd ', other.ID, ')']))
	

#----------------------------------------------------------------------------------------------alignedAnd
    def alignedAnd(self, other, selfReferenceFrame, otherReferenceFrame):
	'''
	performs a frame-by-frame AND of the two movies, 
	lined up with reference to the two specified frames.
	'''

	#shorten the variable names for convenience
	sRefFrame = selfReferenceFrame
	oRefFrame = otherReferenceFrame

	#figure out how many frames we'll be able to use on either side of the reference frame
	framesPostRef = min(self.numBins - sRefFrame, other.numBins - oRefFrame)
	framesPreRef = min(sRefFrame, oRefFrame) 
	#but if the reference frame was 0, we don't want to start comparing at frame -1.
	if(framesPreRef < 0):
	    framesPreRef = 0

	sFirst = sRefFrame - framesPreRef
	oFirst = oRefFrame - framesPreRef
	#now do the and operation aligned to the specified frames.
	temp = numpy.array([map(operator.__and__, self.data.T[sFirst + i], other.data.T[oFirst + i]) for i in range(framesPreRef + framesPostRef)])

	return Movie(temp.T, self.parentSlice, ''.join(['( ', self.ID, '[',str(sFirst), '] AlignedAnd ', other.shortID, '[', str(oFirst), '] )']))
	

#----------------------------------------------------------------------------------------------maxAlignedAnd
    def maxAlignedAnd(self, other):
	'''
	performs a frame-by-frame AND of the two movies
	in this version, we look in each movie for the frame with the most neurons firing
	(in case of tie we take the first frame with the most neurons)
	and we align the two movies to that frame, 
	so the only length gaurantee for the resulting movie is that it's at least one frame.
	'''
	self._checkCompat(other)
	#first thing is to make a vector of length numBins that has 
	#a count of how many neurons were active in each timesetep.
	selfCount =  [sum(self.data[i,j] for i in range(self.numNeurons)) for j in range(self.numBins) ]
	otherCount = [sum(other.data[i,j] for i in range(other.numNeurons)) for j in range(other.numBins) ]

	#now get the reference frames for self (s) and other (o)
	sRefFrame = numpy.argmax(selfCount)
	oRefFrame = numpy.argmax(otherCount)

	#tell us about it
	print('aligning frame '+str(sRefFrame)+' in '+self.ID +'  with frame '+str(oRefFrame)+' in '+other.ID)
	
	return self.alignedAnd(other, sRefFrame, oRefFrame)


#--------------------------------------------------------------------------------------------frameCompare
    def frameCompare(self, other, selfFrameIndex, otherFrameIndex):
        '''
        returns a floating point 'similarity' score 
        for the activity in two frames of two movies

        1 means the frames are identical
        0 means the frames have no neurons in common.
        intermediate values are a function of how many cells
        were active overall so it's more of a percent, but:
        scores closer to zero means the frames are dissimilar
        scores closer to one means the fames are similar
        
        '''
        #we do this by ANDing the two frames, 
        #summing the anded frame
        #and dividing the number of cells active in both (times two)
        #by the sum of the number of cells active in each individually.
        #so a score of one means the frames are identical.
        
        self._checkCompat(other)

	sInd = selfFrameIndex
        oInd = otherFrameIndex

	andies = numpy.array(map(operator.__and__, self.data.T[sInd], other.data.T[oInd]))
	
        asum = sum(andies)
	#the following check avoids division by zero 
        if(asum == 0):
            return float(asum)
        else:
	    numerator = float(asum*2)
	    denominator = float(sum(self.data[:,sInd]) + sum(other.data[:,oInd]))
            return ( numerator / denominator )
        

#----------------------------------------------------------------------------------------bestAlignedAnd
    def bestAlignedAnd(self, other, minNrns=4):
	'''
	aligns the two movies to their most-similar frames with the most number of neurons
	spiking.  Returns a movie of the logical AND based on this alignment.
	'''

	self._checkCompat(other)

	bestFrame = self.findAligningFrame(self, other, minNrns)

	return self.alignedAnd(other, bestFrame[0], bestFrame[1])
	


#--------------------------------------------------------------------------------------findAligningFrame
    def findAligningFrame(self, movie1, movie2, minNrns=4):
	'''
	finds the best matching frame between 2 movies.
	ignores the calling self object - this finds the best frame
	between movie1 and movie2. 
	@param minNrns: the minimum number of neurons to look for in the best frame. default is 4.
	returns: a tuple where the [0] entry is the frame of movie1 and the [1] entry is the frame of movie2.    
	'''
	#step one : only consider frames from each movie with at least minNrns active:
	sInds = []
	oInds = []
	
	for i in xrange(movie1.numBins):
	    if( sum(movie1.data[:,i]) > minNrns):
		sInds.append(i)
	for i in xrange(movie2.numBins):
	    if( sum(movie2.data[:,i]) > minNrns):
		oInds.append(i)

	#now make sure we have anything at all to process.
	if(len(sInds) == 0):
	    print('error: couldnt find any frames with at least %d neurons in %s' %(minNrns, movie1.ID))
	    return None
	if(len(oInds) == 0):
	    print('error: couldnt find any frames with at least %d neurons in %s' %(minNrns, movie2.ID))
	    return None

	#gotta start somewhere, so start at the beginning (a very good place to start)
	bestScore = movie1.frameCompare(movie2, sInds[0], oInds[0])
	bestFrame = (sInds[0], oInds[0])
	#now find the best match frame with the largest number of neurons.
	for sInd in sInds:
	    for oInd in oInds:
		score = movie1.frameCompare(movie2, sInd, oInd)
		if(score > bestScore):
		    bestScore = score
		    bestFrame = (sInd,oInd)
		elif(score == bestScore):
		    #only use these if they have more neurons active, making the score 'more impressive'
		    #but tell us about it so we know it happened.
		    if(sum(movie1.data[:,sInd]) + sum(movie2.data[:,oInd]) > 
		       sum(movie1.data[:,bestFrame[0]]) + sum(movie2.data[:,bestFrame[1]]) ):
			print('%s Frame %d and %s Frame %d, score: %f being replaced in favor of frames %d and %d, which have more neurons active.' %
			      (movie1.ID,bestFrame[0], movie2.ID, bestFrame[1], score, sInd, oInd))
			bestScore = score
			bestFrame = (sInd,oInd)
		    else:
			print('%s Frame %d and %s Frame %d had similarity %f but we\'re not using em.' %
			      (movie1.ID, sInd, movie2.ID, oInd, score))

	print('using %s Frame %d and %s Frame %d , which have similarity: %f' %
	      (movie1.ID, bestFrame[0], movie2.ID, bestFrame[1], bestScore))

	return bestFrame


#------------------------------------------------------------------------------------------------__and__
#unclear why,  but this isn't working. 
#
#    def __and__(self, other):
#	return self.maxAlignedAnd(other)

#-----------------------------------------------------------------------------------------------__sub__
    def __sub__(self, other):
	'''
	returns a movie that is identical to self EXCEPT
	any neuron that fired in any bin of other
	does not spike in the movie we return. 
	
	uses:
	1) looking for multiple ensembles:
	shows the activity of neurons in Self that were
	not recruited by the event in other.
	2) ignoring neurons that are persistently active (not stimulus dependent)    
	'''
	self._checkCompat(other)
	subtracted = numpy.array(self.data, dtype='int')
	#print('sub')

	for i in xrange(self.numNeurons):
	    if( other.timeless[i] ):
		for j in xrange(self.numBins):
		    subtracted[i, j] = 0
	return Movie(subtracted, self.parentSlice, ''.join(['(', self.ID, ' - ', other.shortID, ')']))

#-------------------------------------------------------------------------------------------------timelessAND
    def timelessAnd(self, other):
	'''
	returns a one frame movie (frame0 = timeless)
	that is just the logical AND of the timeless frames
	of the two movies.
	
	uses: 
	1) looking for 'core' neurons that spike in every condition.
	2) looking for multiple ensembles: 
	this can find neurons that always spike under condition A so 
	that you can subtract them from a movie(s) of condition B
	'''
	self._checkCompat(other)
	anded = numpy.zeros(self.numNeurons, dtype='int')
	for i in xrange(self.numNeurons):
	    if( other.timeless[i] and self.timeless[i]):
		anded[i] = numpy.int32(1)

#	print(anded)
		
	#this is annoying, but we need an array of arrays to create a movie,  
	#even tho this is explicitly a one-dimensional manipulation.
	movieProper = []
	for n in anded:
	    movieProper.append([n])

	return Movie(movieProper, self.parentSlice, ''.join(['(', self.ID, ' =TimelessAnd= ', other.shortID, ')']))

#---------------------------------------------------------------------------------------------------_checkCompat
    def _checkCompat(self, other):
	'''
	this just checks that both objects are movies from the same slice...
	called on the first line of most(all) methods for operating on two movies.
	'''
	if(type(self) != type(other)):
	    raise 'Cant compare a Movie with '+str(type(other))
	if self.parentSlice != other.parentSlice:
	    raise 'Cant compare movies from different slices!'


#-------------------------------------------------------------------------------------------------raster
    def raster(self):

	#pylab.figure()
	pylab.matshow(self.data.T) #show a transpose of the data for screensize..
	pylab.show()


#------------------------------------------------------------------------------------------------plotframe
    def plotFrame(self, frameIndex, axes = None, colorString = None):
	''' plots the slice and then colors in the cells active in the given frame.
	NOTE: frames are zero-indexed.
	'''

	fig = None
	if(axes == None):
	    fig = self.parentSlice.plotSlice()
	    axes = fig.gca()
	    
        axes.set_title(self.shortID+' frame: '+str(frameIndex))
	for i in xrange(self.numNeurons):
	    if self.data[i, frameIndex]:
		self.parentSlice.neurons[i].plot(axes, True, colorString)

	if(not (fig == None)): 
	    fig.show()
	
	
	
#-------------------------------------------------------------------------------------------------plotTimeless	
    def plotTimeless(self, axes=None, colorString = 'g', save=False):
	if(axes == None):
	    fig = self.parentSlice.plotSlice()
	    axes = fig.gca()
	
        axes.set_title(self.ID+' timeless')
	for i in xrange(self.numNeurons):
	    if(self.timeless[i]):
		self.parentSlice.neurons[i].plot(axes, True, colorString)

	if(save):
            fig.savefig(''.join([self.parentSlice.directory, self.ID,'-', 'Timless','.png']))
	fig.show()


#--------------------------------------------------------------------------------------------------saveFrames
    def saveFrames(self, show=False):
	'''
	creates a .png file of the activity of each frame of this movie 
	the png files will be in the directory the slice was read in from,
	and will be named the ID of this movie plus 'frame x' where x is the apropriate number.
	frames will be displayed to the screen as well as saved to disk if(show).
	'''
	for i in xrange(self.numBins):
	    f = self.parentSlice.plotSlice()
	    self.plotFrame(i, f.gca())
	    f.savefig(''.join([self.parentSlice.directory, self.ID,'-frame', '%02i'%(i),'.png']))
	    if(show):
		f.show()
	    else:
		pylab.close(f)
	

#---------------------------------------------------------------------------------------saveFrameStrip
    def saveFrameStrip(self, show=False, trim=True):
	'''
	creates a single image which is a (horizontal) series of subplots, one for each frame.
	to economize slightly,  the frameStrip is "trimmed" by default - it starts at the first frame
	with nonzero number of neurons spiking and ends at the last frame with a nonzero number of spikes.
	zero-spike frames in the middle ARE included.
	'''
	pylab.close('all')
	#to start, we need to find the frames to trim:
	startInd = 0
	endInd = self.numBins
	if(trim):
	    #we want the indices to be inclusive at the low end but not the high end,
	    #just like the way we loop. it makes everything good. djikstra wrote an essay about it.
	    while(startInd < endInd and sum(self.data[:,startInd]) == 0):
		startInd += 1
		print('startInd incremented to %d because sum(previous) = %d' % 
		      (startInd, sum(self.data[:,startInd-1])))
	    while(endInd > startInd and sum(self.data[:,endInd-1]) == 0):
		endInd -=1

	numFrames = endInd-startInd
	for i in xrange(numFrames):
	    pylab.subplot(1,numFrames,i+1)
	    pylab.hold(True)
	    axes = pylab.gca()
	    self.parentSlice.plotSlice(axes)
	    self.plotFrame(startInd+i, axes)
	    axes.title.set_size(8)

	#i blame this on matlab, because pylab is just matlab syntax,
	#but there's this thing where it's easier if you don't get handle to the figure
	#until after you've done all the subplot stuff... 
	f = pylab.gcf()
	f.set_size_inches(4*numFrames,3.5)

	f.savefig(''.join([self.parentSlice.directory, self.ID,'.png']))
	if(show):
	    f.show()
	    



#======================================================================================================================
#
#              COMBO MOVIE
#
#======================================================================================================================

class ComboMovie(Movie):
    '''
    this class is intended to help us look at multiple movies from the same slice simultaneously.
    '''
    def __init__(self, movieList, inds=None, alignToInd=0, minNeurons=4):
	'''
	you give it a list of movies and optionally a list of indices. 
	if you don't give indices, then all the movies are used to make the combo movie.
	if you do give indices then only the movies in the list at those indices are used.
	in order to make a combo movie,  we do best-aligned-ands , but we don't do hard AND - we construct an average.
	not all movies are properly aligned by their most similar frames - especially in cases
	of 'echoes' where a very similar pattern appears twice.  for now,  you can control alignment
	by specifying which movie all the others will be aligned against. default is the zeroth movie.
	@param movieList: a list of Movie objects.
	@aram alignToInd: the index, with respect to the movieList, of the movie we should align all the others to.
	@param inds: a list of indices into the movieList of movies that we should consider, if you want to pass
	a whole movielist but only use a subset.
	@param minNeurons: the minimum number of neurons acceptable to align frames against. (gets passed to findAligningFrame)
	'''

        #allow for a copy constructor..
	if movieList == None :
            return
        
	if(inds == None):
	    inds = range(len(movieList))

	self.movieList = []
	self.inds = inds
	self.alignToInd = alignToInd

	#this will be a list of pairs(tuples) of (refMovieFrame, alignedMovieFrame) with one of the two always zero.
	self.alignFrames = []

	refMovie = movieList[inds[alignToInd]]
	self.numBins = refMovie.numBins
	self.parentSlice = refMovie.parentSlice
	self.numNeurons = refMovie.numNeurons

	#make an array of the proper dimensions filled with zeros.
	self.data = numpy.zeros(refMovie.numNeurons * refMovie.numBins).reshape(refMovie.numNeurons, refMovie.numBins)
	

	#make the list of alignFrames to determine how the movies will be lined up for the combo.
	for i in xrange(len(inds)):
	    #add the movie to our list that we keep as part of this object. 
	    self.movieList.append(movieList[inds[i]])
	    print(inds[i])

	    if (i == alignToInd):
		self.alignFrames.append( (0,0) )
	    else:
		bestFrames = self.findAligningFrame(refMovie, movieList[inds[i]], minNrns=minNeurons)
		self.alignFrames.append( (bestFrames[0], bestFrames[1]) )
		

	#combine the movies.
	for i in xrange(len(self.alignFrames)):
	    #shorten the variable names for convenience
	    pair = self.alignFrames[i]
	    sFrame = pair[0]
	    oFrame = pair[1]
	    other = self.movieList[i]

	    #figure out how many frames we'll be able to use on either side of the reference frame
	    framesPostRef = min(self.numBins - sFrame, other.numBins - oFrame)
	    framesPreRef = min(sFrame, oFrame) 

	    sFirst = sFrame - framesPreRef
	    oFirst = oFrame - framesPreRef

	    for f in xrange(framesPreRef + framesPostRef):
		#a frame is a column in the data matrix; 
		#we want to add the other movie to our movie, column by column (frame by frame)
		self.data[:,sFirst + f] += other.data[:,oFirst + f]

	#now normalize (turn it into probabilities!) 
	self.data /= len(self.alignFrames)

	#make a 'timeless' representation, which is.. what probability did each neuron have of firing in this movie
	self.timeless = sum(self.data.T) / self.numBins
	self.ID = self.parentSlice.ID + ' Combo('+''.join(m.shortID+', ' for m in self.movieList)+')'
	self.shortID = 'prob('+''.join(m.shortID+', ' for m in self.movieList)+')'


#---------------------------------------------------------------------------------------------------plotFrame

    def plotFrame(self, frameIndex, axes = None):
	''' plots the slice and then colors in the cells active in the given frame.
	NOTE: frames are zero-indexed.
	'''

	fig = None
	if(axes == None):
	    fig = self.parentSlice.plotSlice()
	    axes = fig.gca()
	    
	
        axes.set_title(self.ID+' frame: '+str(frameIndex))
	for i in xrange(self.numNeurons):
	    if self.data[i, frameIndex]:
		
		#we need to convert this to HTML style hex representation , from 0 to 255, in hex:
		#and what we want is an intense single color if the data is 1.0 in this frame,
		#and grayer/less-intense colors for lower values of data, 
		#and the color should approach white as the value approaches 0.
		val = self.data[i,frameIndex]
		#val is already between 0 and 1. 
		secondaryVal = int( ( 1.0 - val) * 255)
		
		colorStr = '#%02X%02XFF' %(secondaryVal,  secondaryVal)
		

		self.parentSlice.neurons[i].plot(axes, True, colorStr)#str(self.data[i,frameIndex]))

	if(not (fig == None)): 
	    fig.show()



#-----------------------------------------------------------------------------------------------__sub__
    def __sub__(self, other):
	'''
	returns a movie that is identical to self EXCEPT
	any neuron that fired in any bin of other
	does not spike in the movie we return. 
	
	uses:
	1) looking for multiple ensembles:
	shows the activity of neurons in Self that were
	not recruited by the event in other.
	2) ignoring neurons that are persistently active (not stimulus dependent)    
	'''
	self._checkCompat(other)
	subtracted = numpy.array(self.data)

	for i in xrange(self.numNeurons):
	    if( other.timeless[i] ):
		for j in xrange(self.numBins):
		    subtracted[i, j] = 0

        #we want to return a movie that has all the information about the original combo,
	#except it has this subtracted thing too. 
	subtractedMovie = copy.copy(self)
	subtractedMovie.data = subtracted
	subtractedMovie.ID = ''.join(['(',self.ID,' - ', other.ID, ')'])
	subtractedMovie.shortID = ''.join(['(',self.shortID,' - ', other.shortID, ')'])
	return subtractedMovie

#-------------------------------------------------------------------------------------------coreOnly
    def coreOnly(self):
	'''
	returns a (regular) movie with only the cells that were always on in this combo. 
	'''

	return Movie(numpy.array(self.data, dtype='int'), self.parentSlice, ''.join(['Core: ',self.ID]) )


    
